#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Intermediate value模型CPUTrace采集器

Based onQEMU模拟器AndGDB协议采集CPUUse率Trace
"""

import os
import sys
import time
import json
import psutil
import socket
import subprocess
import threading
import numpy as np
from pathlib import Path
from typing import List, Dict, Tuple, Optional
from datetime import datetime


class GDBClient:
    """GDB远程协议客户端"""
    
    def __init__(self, host='localhost', port=1234, timeout=5.0):
        self.host = host
        self.port = port
        self.timeout = timeout
        self.socket = None
        self.connected = False
    
    def connect(self) -> bool:
        """连接ToGDB服务器"""
        try:
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.socket.settimeout(self.timeout)
            self.socket.connect((self.host, self.port))
            self.connected = True
            return True
        except Exception as e:
            print(f"GDB connection failed: {e}")
            return False
    
    def send_command(self, cmd: str) -> str:
        """发送GDBCommand"""
        if not self.connected:
            return ''
        
        # Add校验And
        checksum = sum(ord(c) for c in cmd) % 256
        packet = f"${cmd}#{checksum:02x}"
        
        try:
            self.socket.sendall(packet.encode())
            response = self.socket.recv(4096).decode()
            return response.strip('+$#')
        except Exception as e:
            print(f"GDB command failed: {e}")
            return ''
    
    def continue_execution(self):
        """ContinueExecute"""
        return self.send_command('c')
    
    def close(self):
        """Close连接"""
        if self.socket:
            self.socket.close()
        self.connected = False


class QEMUManager:
    """QEMUProcess管理器"""
    
    def __init__(self, firmware_path: str, qemu_path: str = 'qemu-system-arm'):
        self.firmware_path = firmware_path
        self.qemu_path = qemu_path
        self.process = None
        self.output_thread = None
        self.output_buffer = []
    
    def start(self) -> bool:
        """启动QEMU"""
        if not Path(self.firmware_path).exists():
            print(f"ERROR: Firmware file not found: {self.firmware_path}")
            return False
        
        cmd = [
            self.qemu_path,
            '-M', 'netduinoplus2',
            '-nographic',
            '-semihosting-config', 'enable=on,target=native',
            '-kernel', self.firmware_path
        ]
        
        try:
            print(f"Command: {' '.join(cmd)}")
            self.process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,  # 合并stderrTostdout
                text=True,
                bufsize=1
            )
            
            time.sleep(0.2)  # 短暂WaitingQEMU启动
            
            print(f"QEMU process started (PID: {self.process.pid})")
            return True
            
        except FileNotFoundError:
            print(f"ERROR: QEMU command not found: {self.qemu_path}")
            print("Please install qemu-system-arm")
            return False
        except Exception as e:
            print(f"ERROR: QEMU startup exception: {e}")
            import traceback
            traceback.print_exc()
            return False
    
    def _read_output(self):
        """ReadQEMUOutput（同时ReadstdoutAndstderr）"""
        if not self.process:
            return
        
        # Readstdout
        def read_stdout():
            for line in iter(self.process.stdout.readline, ''):
                if not line:
                    break
                self.output_buffer.append(line.strip())
        
        # Readstderr
        def read_stderr():
            for line in iter(self.process.stderr.readline, ''):
                if not line:
                    break
                self.output_buffer.append(line.strip())
        
        # 启动两个Thread分别Read
        import threading
        t1 = threading.Thread(target=read_stdout, daemon=True)
        t2 = threading.Thread(target=read_stderr, daemon=True)
        t1.start()
        t2.start()
    
    def get_output(self) -> List[str]:
        """获取Output缓冲区"""
        output = self.output_buffer.copy()
        self.output_buffer.clear()
        return output
    
    def stop(self):
        """停止QEMU"""
        if self.process:
            self.process.terminate()
            try:
                self.process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                self.process.kill()
            self.process = None
    
    def is_running(self) -> bool:
        """CheckQEMUYesNoRun"""
        return self.process is not None and self.process.poll() is None


class CPUTraceCollector:
    """CPUUse率Trace采集器"""
    
    def __init__(self, firmware_path: str):
        self.firmware_path = firmware_path
        self.qemu_manager = QEMUManager(firmware_path)
        self.gdb_client = GDBClient()
        self.cpu_monitor_active = False
        self.cpu_samples = []
    
    def start_qemu(self) -> bool:
        """启动QEMU"""
        print("\nStarting QEMU...")
        return self.qemu_manager.start()
    
    def connect_gdb(self) -> bool:
        """连接GDB"""
        print("Connecting GDB...")
        for i in range(3):
            if self.gdb_client.connect():
                print("GDB connected")
                return True
            time.sleep(1)
        print("GDB connection failed")
        return False
    
    def start_cpu_monitoring(self):
        """启动CPUMonitorThread - 持续Run的MonitorLoop"""
        self.cpu_monitor_running = False
        self.cpu_monitor_active = False
        self.cpu_samples = []
        self.qemu_process = None
        
        def monitor():
            # 获取QEMUProcess对象
            try:
                self.qemu_process = psutil.Process(self.qemu_manager.process.pid)
                # InitializeCPUMonitor（First次Call会Return0）
                self.qemu_process.cpu_percent(interval=None)
                time.sleep(0.1)
                print(f"CPU monitor thread ready (PID: {self.qemu_process.pid})")
            except Exception as e:
                print(f"WARNING: Cannot get QEMU process: {e}")
                return
            
            self.cpu_monitor_running = True
            sample_count = 0
            loop_count = 0
            
            # 持续MonitorLoop - No条件高速Sampling
            while self.cpu_monitor_running:
                loop_count += 1
                try:
                    if not self.qemu_process.is_running():
                        print(f"QEMU process stopped (loops: {loop_count})")
                        break
                    
                    cpu = self.qemu_process.cpu_percent(interval=None)
                    
                    # 只HasIn激活时才Save样本
                    if self.cpu_monitor_active:
                        self.cpu_samples.append(cpu)
                        sample_count += 1
                    
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    print(f"CPU monitor: Process no longer exists (loops: {loop_count})")
                    break
                except Exception as e:
                    print(f"CPU monitor error: {e} (loops: {loop_count})")
                    break
            
            print(f"CPU monitor thread stopped (loops: {loop_count}, samples: {sample_count})")
        
        self.monitor_thread = threading.Thread(target=monitor, daemon=True)
        self.monitor_thread.start()
        
        # WaitingMonitor就绪
        timeout = time.time() + 2
        while not self.cpu_monitor_running and time.time() < timeout:
            time.sleep(0.001)
        
        if not self.cpu_monitor_running:
            print("WARNING: CPU monitor thread failed to start")
    
    def start_trace_sampling(self):
        """StartCurrentTrace的CPUSampling"""
        self.cpu_samples = []
        self.cpu_monitor_active = True
        print(f"  Sampling activated (active={self.cpu_monitor_active}, running={self.cpu_monitor_running})")
    
    def stop_trace_sampling(self) -> List[float]:
        """停止CurrentTrace的CPUSampling并Return样本"""
        self.cpu_monitor_active = False
        time.sleep(0.01)  # Waiting更长时间让SamplingCompleted
        samples = self.cpu_samples.copy()
        print(f"  Sampling stopped (collected {len(samples)} samples)")
        return samples
    
    def stop_cpu_monitoring(self):
        """停止CPUMonitor"""
        self.cpu_monitor_active = False
        self.cpu_monitor_running = False
        if hasattr(self, 'monitor_thread'):
            self.monitor_thread.join(timeout=1)
    
    def collect_single_trace(self) -> Dict:
        """采集单条Trace"""
        # 清EmptyCPU样本
        self.cpu_samples.clear()
        
        # WaitingTRACE_START标记
        start_time = time.time()
        trace_started = False
        
        while time.time() - start_time < 10:  # 10秒超时
            output = self.qemu_manager.get_output()
            for line in output:
                if 'TRACE_START' in line:
                    trace_started = True
                    break
            if trace_started:
                break
            time.sleep(0.1)
        
        if not trace_started:
            return {'success': False, 'error': 'TRACE_START超时'}
        
        # WaitingTRACE_END标记
        start_time = time.time()
        trace_ended = False
        intermediate_value = None
        
        while time.time() - start_time < 10:
            output = self.qemu_manager.get_output()
            for line in output:
                if 'TRACE_END' in line:
                    trace_ended = True
                if 'INTERMEDIATE:' in line:
                    try:
                        intermediate_value = int(line.split(':')[1].strip(), 16)
                    except:
                        pass
            if trace_ended:
                break
            time.sleep(0.01)
        
        if not trace_ended:
            return {'success': False, 'error': 'TRACE_END超时'}
        
        # 获取CPU样本
        cpu_trace = np.array(self.cpu_samples.copy())
        
        return {
            'success': True,
            'cpu_trace': cpu_trace,
            'intermediate_value': intermediate_value,
            'trace_length': len(cpu_trace)
        }
    
    def collect_batch(self, num_traces: int) -> Dict:
        """批量采集Trace - 实时ReadQEMUOutput并Based onTRACE_START/TRACE_END控制Sampling"""
        print(f"Reading QEMU output and collecting traces...")
        
        traces = []
        intermediate_values = []
        current_trace_index = None
        
        # 实时ReadQEMU的stdout
        try:
            for line in self.qemu_manager.process.stdout:
                line = line.strip()
                
                # ParseTRACE_START - StartCPUSampling
                if line.startswith('TRACE_START:'):
                    try:
                        index = int(line.split(':')[1])
                        current_trace_index = index
                        # Start新Trace的Sampling窗口
                        self.start_trace_sampling()
                        print(f"[Trace {index}] Start sampling")
                    except:
                        pass
                
                # ParseTRACE_END - 停止CPUSampling并Save
                elif line.startswith('TRACE_END:') and current_trace_index is not None:
                    # 停止Sampling并获取样本
                    cpu_samples = self.stop_trace_sampling()
                    
                    if len(cpu_samples) > 0:
                        trace = np.array(cpu_samples)
                        traces.append(trace)
                        print(f"[Trace {current_trace_index}] End sampling - {len(trace)} samples")
                    else:
                        print(f"[Trace {current_trace_index}] WARNING: No CPU samples")
                    current_trace_index = None
                
                # ParseINTERMEDIATEIntermediate value
                elif line.startswith('INTERMEDIATE:'):
                    try:
                        value = int(line.split(':')[1].strip(), 16)
                        intermediate_values.append(value)
                    except:
                        pass
                
                # 批ProcessCompleted
                elif line == 'BATCH_COMPLETE':
                    print("Batch processing complete")
                    break
        
        except KeyboardInterrupt:
            print("\nInterrupted by user")
        except Exception as e:
            print(f"ERROR during collection: {e}")
        
        success_count = len(traces)
        print(f"\nCollection complete: {success_count} traces, {len(intermediate_values)} intermediate values")
        
        # CheckData质量
        if success_count > 0:
            trace_lengths = [len(t) for t in traces]
            print(f"Trace length stats: min={min(trace_lengths)}, max={max(trace_lengths)}, avg={np.mean(trace_lengths):.1f}")
        
        if success_count == 0:
            print("ERROR: No valid traces collected")
            print("SUGGESTION: Check firmware output for TRACE_START/TRACE_END markers")
        
        return {
            'traces': traces,
            'intermediate_values': intermediate_values,
            'num_traces': len(traces),
            'success_rate': success_count / num_traces if num_traces > 0 else 0
        }
    
    def run_collection(self, num_traces: int) -> Optional[Dict]:
        """RunComplete采集流程"""
        try:
            # 启动QEMU（固件会自动Run）
            if not self.start_qemu():
                return None
            
            # 启动CPUMonitor（MonitorQEMUProcess）
            self.start_cpu_monitoring()
            
            # WaitingExecute并采集Trace
            results = self.collect_batch(num_traces)
            
            if results and results['num_traces'] > 0:
                return results
            else:
                print("ERROR: No valid traces collected")
                return None
            
        except Exception as e:
            print(f"ERROR: Collection failed: {e}")
            import traceback
            traceback.print_exc()
            return None
            
        finally:
            # 清理
            self.stop_cpu_monitoring()
            self.qemu_manager.stop()
    
    def save_results(self, results: Dict, output_path: str, metadata: Dict = None):
        """Save采集Result"""
        # SaveTraceData（NPZFormat）
        npz_path = Path(output_path).with_suffix('.npz')
        
        # 将tracesConvertForobjectArray（SupportNot同长度）
        traces_array = np.array(results['traces'], dtype=object)
        intermediate_values_array = np.array(results['intermediate_values'])
        
        np.savez_compressed(
            npz_path,
            traces=traces_array,
            intermediate_values=intermediate_values_array
        )
        print(f"Trace data saved: {npz_path}")
        
        # Save元Data（JSONFormat）
        json_path = Path(output_path).with_suffix('.json')
        meta = {
            'num_traces': results['num_traces'],
            'success_rate': results['success_rate'],
            'collection_time': datetime.now().isoformat(),
            'firmware_path': self.firmware_path
        }
        if metadata:
            meta.update(metadata)
        
        with open(json_path, 'w') as f:
            json.dump(meta, f, indent=2)
        print(f"Metadata saved: {json_path}")


def main():
    """主Function"""
    import argparse
    
    parser = argparse.ArgumentParser(description='Intermediate value模型CPUTrace采集器')
    parser.add_argument('--firmware', '-f', type=str, required=True, help='固件ELFFilePath')
    parser.add_argument('--traces', '-n', type=int, default=100, help='采集Trace数量')
    parser.add_argument('--output', '-o', type=str, help='Output filePath')
    parser.add_argument('--metadata', '-m', type=str, help='元DataJSONFile')
    
    args = parser.parse_args()
    
    # Read元Data
    metadata = None
    if args.metadata and Path(args.metadata).exists():
        with open(args.metadata, 'r') as f:
            metadata = json.load(f)
    
    # Create采集器
    collector = CPUTraceCollector(args.firmware)
    
    # Run采集
    results = collector.run_collection(args.traces)
    
    if results and results['num_traces'] > 0:
        # SaveResult
        if args.output:
            collector.save_results(results, args.output, metadata)
        else:
            # DefaultOutputPath
            output_path = f"traces_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
            collector.save_results(results, output_path, metadata)
        
        print(f"\nCollection completed successfully!")
        print(f"   Success rate: {results['success_rate']:.1%}")
        print(f"   Traces: {results['num_traces']}")
    else:
        print(f"\nERROR: Collection failed")


if __name__ == '__main__':
    main()
